//
//  H264HWDecoder.m
//  IntellDev_new
//
//  Created by os on 2021/3/11.
//

#import "H264HWDecoder.h"

@implementation H264HWDecoder

- (void)dealloc {
    if (self.decompressionSession) {
        VTDecompressionSessionInvalidate(self.decompressionSession);
        CFRelease(self.decompressionSession);
        self.decompressionSession = NULL;
    }
    if (self.formatDesc) {
        CFRelease(self.formatDesc);
        self.formatDesc = NULL;
    }
}

- (int)DecodeH264Frames:(unsigned char *)inputBuffer withLength:(int)length {
    OSStatus status = -1;
    
    uint8_t *data = NULL;
    uint8_t *sps = NULL;
    uint8_t *pps = NULL;
    
    int startCodeIndex = 0;
    int secondStartCodeIndex = 0;
    int thirdStartCodeIndex = 0;
    
    long blockLength = 0;
    
    CMSampleBufferRef sampleBuffer = NULL;
    CMBlockBufferRef blockBuffer = NULL;
    
    int nalu_type = (inputBuffer[startCodeIndex + 4] & 0x1F);
    
    if (nalu_type != 7 && _formatDesc == NULL) {
        NSLog(@"Video error: Frame is not an I Frame and format description is null");
        return -1;
    }
    
    if (nalu_type == 7) {
        for (int i = startCodeIndex + 4; i < startCodeIndex + 256; i++) {
            if (inputBuffer[i] == 0x00 && inputBuffer[i+1] == 0x00 && inputBuffer[i+2] == 0x00 && inputBuffer[i+3] == 0x01) {
                secondStartCodeIndex = i;
                _spsSize = secondStartCodeIndex;
                break;
            }
        }
        nalu_type = (inputBuffer[secondStartCodeIndex + 4] & 0x1F);
    }
    
    if (nalu_type == 8) {
        for (int i = _spsSize + 4; i < _spsSize + 128; i++) {
            if (inputBuffer[i] == 0x00 && inputBuffer[i+1] == 0x00 && inputBuffer[i+2] == 0x01) {
                thirdStartCodeIndex = i;
                _ppsSize = thirdStartCodeIndex - _spsSize;
                break;
            }
        }
        
        sps = malloc(_spsSize - 4);
        pps = malloc(_ppsSize - 4);
        
        memcpy(sps, &inputBuffer[4], _spsSize-4);
        memcpy(pps, &inputBuffer[_spsSize+4], _ppsSize-4);
        
        uint8_t *parameterSetPointers[2] = {sps, pps};
        size_t parameterSetSizes[2] = {_spsSize-4, _ppsSize-4};
        
        status = CMVideoFormatDescriptionCreateFromH264ParameterSets(kCFAllocatorDefault, 2, (const uint8_t * const *)parameterSetPointers, parameterSetSizes, 4, &_formatDesc);
        
        if (status != noErr) {
            NSLog(@"CMVideoFormatDescriptionCreateFromH264ParameterSets Error type: %d", status);
        }
        
        nalu_type = (inputBuffer[thirdStartCodeIndex + 3] & 0x1F);
    }
    
    if (nalu_type == 5) {
        int offset = _spsSize + _ppsSize;
        blockLength = length - offset;
        data = malloc(blockLength);
        data = memcpy(data, &inputBuffer[offset], blockLength);
        
        uint32_t dataLength32 = htonl(blockLength-4);
        memcpy(data, &dataLength32, sizeof(uint32_t));
        
        status = CMBlockBufferCreateWithMemoryBlock(NULL, data, blockLength, kCFAllocatorNull, NULL, 0, blockLength, 0, &blockBuffer);
        
        if (status != noErr) {
            NSLog(@"I Frame: CMBlockBufferCreateWithMemoryBlock Error type: %d", status);
        }
    }
    
    if (nalu_type == 1) {
        blockLength = length;
        data = malloc(blockLength);
        data = memcpy(data, &inputBuffer[0], blockLength);
        
        uint32_t dataLength32 = htonl(blockLength-4);
        memcpy(data, &dataLength32, sizeof(uint32_t));
        
        status = CMBlockBufferCreateWithMemoryBlock(NULL, data, blockLength, kCFAllocatorNull, NULL, 0, blockLength, 0, &blockBuffer);
        
        if (status != noErr) {
            NSLog(@"P Frame: CMBlockBufferCreateWithMemoryBlock Error type: %d", status);
        }
    }
    
    if (status == noErr) {
        const size_t sampleSize = blockLength;
        status = CMSampleBufferCreate(kCFAllocatorDefault, blockBuffer, true, NULL, NULL, _formatDesc, 1, 0, NULL, 1, &sampleSize, &sampleBuffer);
        
        if (status != noErr) {
            NSLog(@"CMSampleBufferCreate Error type: %d", status);
        }
    }
    
    if (status == noErr) {
        CFArrayRef attachments = CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, YES);
        CFMutableDictionaryRef dict = (CFMutableDictionaryRef)CFArrayGetValueAtIndex(attachments, 0);
        CFDictionarySetValue(dict, kCMSampleAttachmentKey_DisplayImmediately, kCFBooleanTrue);
        
        if ([self.updateDelegate respondsToSelector:@selector(updateDecodedSampleBuffer:)]) {
            [self.updateDelegate updateDecodedSampleBuffer:sampleBuffer];
        }
        
    }
    
    if (data) {
        free(data);
        data = NULL;
    }
    if (sps) {
        free(sps);
        sps = NULL;
    }
    if (pps) {
        free(pps);
        pps = NULL;
    }
    
    return 0;
}



@end
